package top.kangyaocoding.tech.common.utils;

import lombok.extern.slf4j.Slf4j;
import org.redisson.api.RList;
import org.springframework.ai.chat.memory.ChatMemory;
import org.springframework.ai.chat.messages.AssistantMessage;
import org.springframework.ai.chat.messages.Message;
import org.springframework.ai.chat.messages.SystemMessage;
import org.springframework.ai.chat.messages.UserMessage;
import org.springframework.stereotype.Component;
import top.kangyaocoding.tech.common.constants.Constants;
import top.kangyaocoding.tech.common.redis.IRedisService;
import top.kangyaocoding.tech.model.ChatEntity;
import top.kangyaocoding.tech.model.MessageTypeEnum;

import java.time.Duration;
import java.util.ArrayList;
import java.util.Collections;
import java.util.List;

/**
 * 描述: 重写 ChatMemory 实现持久化到 Redis
 *
 * @author K·Herbert
 * @since 2025-04-16 11:58
 */
@Slf4j
@Component
public class ChatRedisMemory implements ChatMemory {
    private static final String CHAT_MEMORY_KEY_PREFIX = Constants.CHAT_MEMORY_KEY_PREFIX;
    private final IRedisService redissonService;

    public ChatRedisMemory(IRedisService redissonService) {
        this.redissonService = redissonService;
    }

    @Override
    public void add(String conversationId, List<Message> messages) {
        String key = CHAT_MEMORY_KEY_PREFIX + conversationId;
        RList<ChatEntity> chatList = redissonService.getList(key);

        List<ChatEntity> chatEntities = messages.stream()
                .map(msg -> ChatEntity.builder()
                        .chatId(conversationId)
                        .type(msg.getMessageType().getValue())
                        .text(msg.getText())
                        .build())
                .toList();
        chatList.addAll(chatEntities);
        chatList.expire(Duration.ofMinutes(30));
    }

    @Override
    public List<Message> get(String conversationId, int lastN) {
        String key = CHAT_MEMORY_KEY_PREFIX + conversationId;
        RList<ChatEntity> list = redissonService.getList(key);
        int size = list.size();

        if (size == 0) {
            return Collections.emptyList();
        }

        int start = Math.max(0, size - lastN);
        List<ChatEntity> historyChats = list.range(start, size - 1);

        List<Message> results = new ArrayList<>();
        for (ChatEntity chat : historyChats) {
            try {
                MessageTypeEnum type = MessageTypeEnum.fromValue(chat.getType());
                switch (type) {
                    case USER:
                        results.add(new UserMessage(chat.getText()));
                        break;
                    case ASSISTANT:
                        results.add(new AssistantMessage(chat.getText()));
                        break;
                    case SYSTEM:
                        results.add(new SystemMessage(chat.getText()));
                        break;
                    default:
                        break;
                }
            } catch (IllegalArgumentException e) {
                // 未知的消息类型
                log.error("Unknown message type: {}", chat.getType(), e);
            }
        }
        return results;
    }

    @Override
    public void clear(String conversationId) {
        String key = CHAT_MEMORY_KEY_PREFIX + conversationId;
        redissonService.remove(key);
    }
}

